/*
 *  Copyright (c) 2011 Lukáš Karas <lukas.karas@centrum.cz>
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License version 2 as
 *  published by the Free Software Foundation;
 *
 *  Software distributed under the License is distributed on an "AS IS"
 *  basis, WITHOUT WARRANTY OF ANY KIND, either express or implied.
 *
 *  See the GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program; if not, write to the Free Software
 *  Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
 */

package cz.karry.jtw;

import com.sun.tools.attach.VirtualMachine;
import com.sun.tools.attach.VirtualMachineDescriptor;
import java.io.BufferedOutputStream;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.lang.management.ManagementFactory;
import java.lang.management.ThreadInfo;
import java.lang.management.ThreadMXBean;
import java.text.DateFormat;
import java.text.SimpleDateFormat;
import java.util.Date;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import javax.management.MBeanServerConnection;
import javax.management.remote.JMXConnector;
import javax.management.remote.JMXConnectorFactory;
import javax.management.remote.JMXServiceURL;
import org.apache.log4j.Logger;

/**
 *
 * @author karry
 */
public class ThreadWatchDog extends AbstractPoolingThread {

  private static final Logger logger = Logger.getLogger(ThreadWatchDog.class);
  private static final DateFormat dateFormat = new SimpleDateFormat("dd.MM.yyyy HH:mm:ss");
  private static final String delimiter = "\n===================================================================\n";
  private final VirtualMachineDescriptor machine;
  private final String statDir;
  private OutputStreamWriter output;
  private static final String CONNECTOR_ADDRESS =
          "com.sun.management.jmxremote.localConnectorAddress";
  private ThreadMXBean threadBean;
  private final long warningTime;
  private final Map<Long, BlockedThreadInfo> blockedThreadsMap = new HashMap<Long, BlockedThreadInfo>();

  ThreadWatchDog(VirtualMachineDescriptor m, String statDir, long warningTime) {
    this.machine = m;
    this.warningTime = warningTime;
    this.statDir = statDir;
    this.name = "threadPool-" + machine.id();
  }

  @Override
  public boolean start(long poolTime) {
    try {
      // setup VM
      JMXServiceURL jmxUrl = getURLForPid(machine.id());
      JMXConnector jmxc = JMXConnectorFactory.connect(jmxUrl);
      MBeanServerConnection server = jmxc.getMBeanServerConnection();

      threadBean = ManagementFactory.newPlatformMXBeanProxy(server, ManagementFactory.THREAD_MXBEAN_NAME, ThreadMXBean.class);

      if (!threadBean.isThreadCpuTimeSupported()) {
        throw new Exception("This VM does not support thread CPU time monitoring");
      } else {
        threadBean.setThreadCpuTimeEnabled(true);
      }
      if (!threadBean.isThreadContentionMonitoringSupported()) {
        throw new Exception("This VM does not support thread contention monitoring");
      } else {
        threadBean.setThreadContentionMonitoringEnabled(true);
      }
      if (!threadBean.isObjectMonitorUsageSupported()) {
        throw new Exception("This VM does not supports monitoring of object monitor usage");
      }
      if (!threadBean.isSynchronizerUsageSupported()) {
        throw new Exception("This VM does not supports monitoring of ownable synchronizer usage");
      }

      // create and open stat file
      File f = new File(statDir, machine.id());
      int i = 0;
      while (f.exists()) {
        f = new File(statDir, machine.id() + "." + i);
        i++;
      }
      output = new OutputStreamWriter(new BufferedOutputStream(new FileOutputStream(f)));
      output.write("watchdog for process \"" + (machine.id() + "/" + machine.displayName()) + "\" started at " + dateFormat.format(new Date()) + "\n\n");

      // all is ok, start pooling
      super.start(poolTime);
      logger.info("watchdog started for \"" + (machine.id() + "/" + machine.displayName()) + "\"; stat file " + f.getPath());
    } catch (Exception ex) {
      logger.error("Exception while starting process watchdog for " + machine.id() + "/" + machine.displayName(), ex);
    }
    return true;
  }

  @Override
  public boolean stop() {
    if (!super.stop())
      return false;
    logger.info("watchdog stopped for " + (machine.id() + "/" + machine.displayName()));

    try {
      if (output != null) {
        output.flush();
        output.close();
      }
    } catch (IOException ex) {
      logger.error("exception while closing output", ex);
    }
    return true;
  }

  @Override
  void doJob() {

    long[] tids = null;
    ThreadInfo[] tinfos = null;
    try {
      tids = threadBean.getAllThreadIds();
      tinfos = threadBean.getThreadInfo(tids, true, true);
    } catch (java.lang.reflect.UndeclaredThrowableException ex) {
      // TODO: detect disconnected VM better
      if (ex.getCause() != null && ex.getCause() instanceof java.rmi.ConnectException) {
        logger.warn("remote VM was probably stopped");
        stop();
      } else {
        logger.error("undeclared exception", ex);
      }
      return;
    }

    Set<Long> blockedTids = new HashSet<Long>();
    long tid;
    BlockedThreadInfo blockInfo;
    boolean flushOutput = false;

    for (ThreadInfo info : tinfos) {
      // iterate all blocked threads...
      if (info != null && info.getThreadState() == Thread.State.BLOCKED) {
        tid = info.getThreadId();
        blockedTids.add(tid);
        blockInfo = blockedThreadsMap.get(tid);
        if (blockInfo == null) {  // if this thread is newly blocked, add to map...
          blockInfo = new BlockedThreadInfo(info);
          blockedThreadsMap.put(tid, blockInfo);
        } else {
          if (blockInfo.getBlockedCount() != info.getBlockedCount()) { // this thrad is blocked again - update info...
            blockInfo.setBlockedCount(info.getBlockedCount());
            blockInfo.setBlockedTime(info.getBlockedTime());
            blockInfo.setWarned(false);
          } else { // this is blocked still (from previous invocation)
            if (info.getBlockedTime() - blockInfo.getBlockedTime() >= warningTime && (!blockInfo.warnedAlready())) { // it is blocked too long, write warning to stat
              try {
                if (!flushOutput)
                  output.write(delimiter);

                writeThreadWarning(info, tinfos, blockInfo.getBlockedTime());
              } catch (IOException ex) {
                logger.error("exception while store blocked thread stat", ex);
              } finally {
                flushOutput = true;
                blockInfo.setWarned(true);
              }
            }
          }
        }
      }
    }

    // remove unblocked threads from map
    Set<Long> copy = new HashSet<Long>(blockedThreadsMap.keySet());
    for (Long id : copy) {
      if (!blockedTids.contains(id))
        blockedThreadsMap.remove(id);
    }

    if (flushOutput) {
      try {
        output.flush();
      } catch (IOException ex) {
        logger.error("exception while flushing output", ex);
      }
    }
  }

  private JMXServiceURL getURLForPid(String pid) throws Exception {

    // attach to the target application
    final VirtualMachine vm = VirtualMachine.attach(pid);

    // get the connector address
    String connectorAddress = vm.getAgentProperties().getProperty(CONNECTOR_ADDRESS);

    // no connector address, so we start the JMX agent
    if (connectorAddress == null) {
      String agent = vm.getSystemProperties().getProperty("java.home")
              + File.separator + "lib" + File.separator + "management-agent.jar";
      vm.loadAgent(agent);

      // agent is started, get the connector address
      connectorAddress = vm.getAgentProperties().getProperty(CONNECTOR_ADDRESS);
      assert connectorAddress != null;
    }
    return new JMXServiceURL(connectorAddress);
  }

  private void writeThreadWarning(ThreadInfo info, ThreadInfo[] tinfos, long lastBlockStart) throws IOException {
    String msg = dateFormat.format(new Date()) + " - thread " + info.getThreadId() + " \"" + info.getThreadName() + "\" is blocked more than " + (info.getBlockedTime() - lastBlockStart) + "ms!";
    logger.debug(msg);

    output.write("\n" + msg + "\n");
    output.write(info.toString() + "");
    output.write("blocked by:\n");

    long blockedBy = info.getLockOwnerId();
    for (ThreadInfo i : tinfos) {
      if (i!= null && i.getThreadId() == blockedBy) {
        output.write(i.toString());
        break;
      }
    }
  }
}
